from neo4j import GraphDatabase, exceptions
import json
import os
import sys
import time
from typing import List, Dict, Any
from datetime import datetime

# 获取项目根目录
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
PROCESSED_DIR = os.path.join(PROJECT_ROOT, "processed")


class KnowledgeGraphBuilder:
    def __init__(self, uri: str, user: str, password: str, database: str = "neo4j"):
        """初始化Neo4j连接"""
        self.uri = uri
        self.user = user
        self.database = database
        self.connection_successful = False

        try:
            self.driver = GraphDatabase.driver(
                uri,
                auth=(user, password),
                max_connection_lifetime=3600,
                connection_timeout=30,
                max_connection_pool_size=50
            )
            self.verify_connection()
            self.connection_successful = True
        except Exception as e:
            raise Exception(f"Neo4j连接初始化失败: {e}")

    def verify_connection(self) -> None:
        """验证数据库连接"""
        try:
            with self.driver.session(database=self.database) as session:
                result = session.run("RETURN 1 AS test")
                record = result.single()
                if record and record["test"] == 1:
                    print("✓ Neo4j连接成功")
                else:
                    raise Exception("连接测试失败")
        except exceptions.ServiceUnavailable:
            raise Exception("Neo4j服务未启动，请检查服务状态")
        except exceptions.AuthError:
            raise Exception("Neo4j认证失败，请检查账号密码")
        except Exception as e:
            raise Exception(f"连接验证失败: {e}")

    def check_neo4j_version(self) -> str:
        """检查Neo4j版本"""
        try:
            with self.driver.session(database=self.database) as session:
                result = session.run(
                    "CALL dbms.components() YIELD versions, edition RETURN versions[0] AS version, edition")
                record = result.single()
                if record:
                    version = record["version"]
                    edition = record["edition"]
                    print(f"✓ Neo4j版本: {version} ({edition} edition)")
                    return version
                return "未知"
        except Exception as e:
            print(f"⚠ 获取版本信息失败: {e}")
            return "未知"

    def create_constraints(self) -> None:
        """创建所有必要的约束"""
        constraints = [
            {
                "name": "entity_name_unique",
                "query": """
                    CREATE CONSTRAINT entity_name_unique IF NOT EXISTS
                    FOR (e:Entity) REQUIRE e.name IS UNIQUE
                """,
                "fallback_query": """
                    CREATE CONSTRAINT ON (e:Entity) ASSERT e.name IS UNIQUE
                """
            },
            {
                "name": "relation_type_index",
                "query": """
                    CREATE INDEX relation_type_index IF NOT EXISTS
                    FOR ()-[r:RELATION]-() ON (r.type)
                """,
                "fallback_query": """
                    CREATE INDEX ON :RELATION(type)
                """
            }
        ]

        for constraint in constraints:
            try:
                with self.driver.session(database=self.database) as session:
                    # 尝试新语法
                    try:
                        result = session.run(constraint["query"])
                        result.consume()
                        print(f"✓ 约束 '{constraint['name']}' 创建成功")
                    except exceptions.ClientError:
                        # 尝试回退语法
                        try:
                            result = session.run(constraint["fallback_query"])
                            result.consume()
                            print(f"✓ 约束 '{constraint['name']}' 创建成功（使用回退语法）")
                        except exceptions.ClientError as e:
                            if "already exists" in str(e).lower():
                                print(f"✓ 约束 '{constraint['name']}' 已存在")
                            else:
                                print(f"⚠ 约束 '{constraint['name']}' 创建失败: {e}")
            except Exception as e:
                print(f"⚠ 约束 '{constraint['name']}' 创建错误: {e}")

    def import_entities(self, entities: List[str]) -> Dict[str, Any]:
        """导入实体（节点）"""
        if not entities:
            print("⚠ 没有实体数据可导入")
            return {"success": 0, "failed": 0, "total": 0}

        total = len(entities)
        success_count = 0
        failed_count = 0
        start_time = time.time()

        print(f"📊 开始导入 {total} 个实体...")

        try:
            with self.driver.session(database=self.database) as session:
                batch_size = 500

                for i in range(0, total, batch_size):
                    batch = entities[i:i + batch_size]

                    try:
                        result = session.run("""
                            UNWIND $entities AS entity
                            MERGE (e:Entity {name: entity})
                            SET e.created_at = datetime()
                            RETURN count(e) as count
                        """, entities=batch)

                        record = result.single()
                        if record:
                            success_count += record["count"]

                        # 显示进度
                        processed = min(i + batch_size, total)
                        progress = (processed / total) * 100
                        print(f"📈 实体导入进度: {processed}/{total} ({progress:.1f}%)")

                    except Exception as e:
                        print(f"⚠ 批量导入失败: {e}")
                        failed_count += len(batch)

                    # 小延迟避免过载
                    time.sleep(0.1)

        except Exception as e:
            print(f"❌ 实体导入失败: {e}")
            failed_count = total - success_count

        elapsed_time = time.time() - start_time
        print(f"✅ 实体导入完成 - 成功: {success_count}, 失败: {failed_count}, 耗时: {elapsed_time:.2f}秒")

        return {
            "success": success_count,
            "failed": failed_count,
            "total": total,
            "time": elapsed_time
        }

    def import_relations(self, relations: List[Dict]) -> Dict[str, Any]:
        """导入关系（边）"""
        if not relations:
            print("⚠ 没有关系数据可导入")
            return {"success": 0, "failed": 0, "total": 0}

        total = len(relations)
        success_count = 0
        failed_count = 0
        start_time = time.time()

        print(f"📊 开始导入 {total} 个关系...")

        try:
            with self.driver.session(database=self.database) as session:
                batch_size = 300

                for i in range(0, total, batch_size):
                    batch = relations[i:i + batch_size]

                    try:
                        result = session.run("""
                            UNWIND $relations AS rel
                            MATCH (s:Entity {name: rel.source}), (t:Entity {name: rel.target})
                            MERGE (s)-[r:RELATION {type: rel.type}]->(t)
                            SET r.text = rel.text,
                                r.created_at = datetime()
                            RETURN count(r) as count
                        """, relations=batch)

                        record = result.single()
                        if record:
                            success_count += record["count"]

                        # 显示进度
                        processed = min(i + batch_size, total)
                        progress = (processed / total) * 100
                        print(f"📈 关系导入进度: {processed}/{total} ({progress:.1f}%)")

                    except Exception as e:
                        print(f"⚠ 批量关系导入失败: {e}")
                        failed_count += len(batch)

                    # 小延迟避免过载
                    time.sleep(0.1)

        except Exception as e:
            print(f"❌ 关系导入失败: {e}")
            failed_count = total - success_count

        elapsed_time = time.time() - start_time
        print(f"✅ 关系导入完成 - 成功: {success_count}, 失败: {failed_count}, 耗时: {elapsed_time:.2f}秒")

        return {
            "success": success_count,
            "failed": failed_count,
            "total": total,
            "time": elapsed_time
        }

    def verify_import(self) -> Dict[str, Any]:
        """验证导入结果"""
        try:
            with self.driver.session(database=self.database) as session:
                # 统计实体数量
                entity_result = session.run("MATCH (e:Entity) RETURN count(e) as entity_count")
                entity_count = entity_result.single()["entity_count"] if entity_result else 0

                # 统计关系数量
                relation_result = session.run("MATCH ()-[r:RELATION]->() RETURN count(r) as relation_count")
                relation_count = relation_result.single()["relation_count"] if relation_result else 0

                # 统计关系类型
                type_result = session.run(
                    "MATCH ()-[r:RELATION]->() RETURN r.type as type, count(r) as count ORDER BY count DESC")
                relation_types = []
                for record in type_result:
                    relation_types.append({
                        "type": record["type"],
                        "count": record["count"]
                    })

                return {
                    "entities": entity_count,
                    "relations": relation_count,
                    "relation_types": relation_types
                }

        except Exception as e:
            print(f"⚠ 验证导入结果失败: {e}")
            return {"entities": 0, "relations": 0, "relation_types": []}

    def close(self) -> None:
        """关闭数据库连接"""
        if hasattr(self, 'driver'):
            try:
                self.driver.close()
                print("✓ Neo4j连接已关闭")
            except Exception as e:
                print(f"⚠ 关闭连接时出错: {e}")


def load_data(file_path: str, default_data=None):
    """加载JSON数据文件"""
    if default_data is None:
        default_data = []

    if not os.path.exists(file_path):
        print(f"⚠ 文件 {file_path} 不存在")
        return default_data

    try:
        with open(file_path, "r", encoding="utf-8") as f:
            data = json.load(f)
            print(f"✓ 加载文件: {file_path} ({len(data) if isinstance(data, list) else 'dict'} 条数据)")
            return data
    except json.JSONDecodeError as e:
        print(f"❌ 文件格式错误: {file_path} - {e}")
        return default_data
    except Exception as e:
        print(f"❌ 读取文件失败: {file_path} - {e}")
        return default_data


def test_neo4j_connection(uri: str, user: str, password: str, database: str = "neo4j") -> bool:
    """测试Neo4j连接"""
    try:
        driver = GraphDatabase.driver(uri, auth=(user, password))
        with driver.session(database=database) as session:
            result = session.run("RETURN 1 AS test")
            record = result.single()
            success = record and record["test"] == 1
            if success:
                print("✓ Neo4j连接测试成功")
            else:
                print("✗ Neo4j连接测试失败")
            return success
    except exceptions.ServiceUnavailable:
        print("✗ Neo4j服务未启动，请启动Neo4j服务")
        return False
    except exceptions.AuthError:
        print("✗ Neo4j认证失败，请检查用户名和密码")
        return False
    except Exception as e:
        print(f"✗ Neo4j连接错误: {e}")
        return False
    finally:
        if 'driver' in locals():
            driver.close()


def main():
    """主函数"""
    print("=" * 60)
    print("🧠 知识图谱构建工具")
    print("=" * 60)

    print(f"📁 项目根目录: {PROJECT_ROOT}")
    print(f"📂 处理数据目录: {PROCESSED_DIR}")

    # Neo4j连接配置
    NEO4J_CONFIG = {
        "uri": "bolt://localhost:7687",
        "user": "neo4j",
        "password": "123456",  # 替换为你的实际密码
        "database": "neo4j"
    }

    print(f"\n🔗 Neo4j配置:")
    print(f"   URI: {NEO4J_CONFIG['uri']}")
    print(f"   User: {NEO4J_CONFIG['user']}")
    print(f"   Password: {'*' * len(NEO4J_CONFIG['password']) if NEO4J_CONFIG['password'] else '未设置'}")
    print(f"   Database: {NEO4J_CONFIG['database']}")

    # 测试连接
    print(f"\n🔍 测试Neo4j连接...")
    if not test_neo4j_connection(**NEO4J_CONFIG):
        print("\n❌ 连接失败，请检查:")
        print("   1. Neo4j服务是否启动 (neo4j start)")
        print("   2. 连接URI是否正确")
        print("   3. 用户名和密码是否正确")
        print("   4. 防火墙是否阻止了7687端口")
        print("   5. Neo4j浏览器是否可访问 (http://localhost:7474)")
        return False

    # 检查数据目录
    if not os.path.exists(PROCESSED_DIR):
        print(f"\n❌ 目录 {PROCESSED_DIR} 不存在")
        print("   请先运行 data_processor.py 处理数据")
        return False

    # 加载数据
    entities_path = os.path.join(PROCESSED_DIR, "entities.json")
    relations_path = os.path.join(PROCESSED_DIR, "relations.json")

    print(f"\n📊 加载数据文件...")
    entities = load_data(entities_path, [])
    relations = load_data(relations_path, [])

    if not entities and not relations:
        print("❌ 没有找到任何实体或关系数据")
        print("   请先运行 data_processor.py 处理数据")
        return False

    print(f"   ✅ 加载 {len(entities)} 个实体")
    print(f"   ✅ 加载 {len(relations)} 个关系")

    # 显示示例数据
    if entities:
        print(f"\n📋 实体示例:")
        for i, entity in enumerate(entities[:3]):
            print(f"   {i + 1}. {entity}")

    if relations:
        print(f"\n📋 关系示例:")
        for i, relation in enumerate(relations[:2]):
            print(f"   {i + 1}. {relation['source']} → {relation['target']} ({relation['type']})")

    # 构建知识图谱
    print(f"\n🚀 开始构建知识图谱...")

    try:
        kg_builder = KnowledgeGraphBuilder(**NEO4J_CONFIG)

        if not kg_builder.connection_successful:
            return False

        # 检查版本
        kg_builder.check_neo4j_version()

        # 创建约束
        print(f"\n🔒 创建数据库约束...")
        kg_builder.create_constraints()

        # 导入数据
        print(f"\n📥 导入数据到Neo4j...")
        entity_result = kg_builder.import_entities(entities)
        relation_result = kg_builder.import_relations(relations)

        # 验证导入结果
        print(f"\n🔍 验证导入结果...")
        verification = kg_builder.verify_import()

        print(f"\n🎉 知识图谱构建完成!")
        print("=" * 40)
        print(f"📊 最终统计:")
        print(f"   • 实体总数: {verification['entities']}")
        print(f"   • 关系总数: {verification['relations']}")
        print(f"   • 关系类型: {len(verification['relation_types'])} 种")

        if verification['relation_types']:
            print(f"\n🔗 关系类型统计:")
            for i, rel_type in enumerate(verification['relation_types'][:5]):
                print(f"   {i + 1}. {rel_type['type']}: {rel_type['count']} 个")
            if len(verification['relation_types']) > 5:
                print(f"   ... 还有 {len(verification['relation_types']) - 5} 种类型")

        print(f"\n⏱ 性能统计:")
        print(f"   • 实体导入: {entity_result['success']}/{entity_result['total']} ({entity_result['time']:.2f}s)")
        print(
            f"   • 关系导入: {relation_result['success']}/{relation_result['total']} ({relation_result['time']:.2f}s)")

        print(f"\n💡 下一步:")
        print(f"   1. 打开 Neo4j Browser: http://localhost:7474")
        print(f"   2. 运行查询: MATCH (n:Entity) RETURN n LIMIT 25")
        print(f"   3. 运行查询: MATCH p=()-[r:RELATION]->() RETURN p LIMIT 25")

        return True

    except Exception as e:
        print(f"\n❌ 构建知识图谱失败: {e}")
        import traceback
        traceback.print_exc()
        return False

    finally:
        if 'kg_builder' in locals():
            kg_builder.close()


if __name__ == "__main__":
    # 记录开始时间
    start_time = datetime.now()
    print(f"⏰ 开始时间: {start_time.strftime('%Y-%m-%d %H:%M:%S')}")

    # 运行主程序
    success = main()

    # 记录结束时间
    end_time = datetime.now()
    duration = end_time - start_time
    print(f"\n⏰ 结束时间: {end_time.strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"⏱ 总耗时: {duration.total_seconds():.2f} 秒")

    # 退出